/**
 * Copyright (c) 2010 Yahoo! Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License. You
 * may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License. See accompanying
 * LICENSE file.
 */

package com.impetus.kundera.ycsb.benchmark;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.Vector;

import org.apache.cassandra.thrift.AuthenticationRequest;
import org.apache.cassandra.thrift.Cassandra;
import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnOrSuperColumn;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.ColumnPath;
import org.apache.cassandra.thrift.ConsistencyLevel;
import org.apache.cassandra.thrift.KeyRange;
import org.apache.cassandra.thrift.KeySlice;
import org.apache.cassandra.thrift.Mutation;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.scale7.cassandra.pelops.Bytes;

import com.yahoo.ycsb.ByteArrayByteIterator;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.StringByteIterator;
import common.Logger;

//XXXX if we do replication, fix the consistency levels
/**
 * Cassandra 1.0.6 client for YCSB framework
 */
public class ThriftClient extends DB
{
    private static Logger logger = Logger.getLogger(ThriftClient.class);

    static Random random = new Random();

    public static final int Ok = 0;

    public static final int Error = -1;

    public static final ByteBuffer emptyByteBuffer = ByteBuffer.wrap(new byte[0]);

    public int ConnectionRetries;

    public String column_family;

    public static final String CONNECTION_RETRY_PROPERTY = "cassandra.connectionretries";

    public static final String CONNECTION_RETRY_PROPERTY_DEFAULT = "300";

    public static final String USERNAME_PROPERTY = "cassandra.username";

    public static final String PASSWORD_PROPERTY = "cassandra.password";

    public static final String COLUMN_FAMILY_PROPERTY = "columnfamilyOrTable";

    // public static final String COLUMN_FAMILY_PROPERTY_DEFAULT = "data";

    public static final String READ_CONSISTENCY_LEVEL_PROPERTY = "cassandra.readconsistencylevel";

    public static final String READ_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";

    public static final String WRITE_CONSISTENCY_LEVEL_PROPERTY = "cassandra.writeconsistencylevel";

    public static final String WRITE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";

    public static final String SCAN_CONSISTENCY_LEVEL_PROPERTY = "cassandra.scanconsistencylevel";

    public static final String SCAN_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";

    public static final String DELETE_CONSISTENCY_LEVEL_PROPERTY = "cassandra.deleteconsistencylevel";

    public static final String DELETE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT = "ONE";

    TTransport tr;

    Cassandra.Client client;

    boolean _debug = false;

    String _table = "";

    Exception errorexception = null;

    List<Mutation> mutations = new ArrayList<Mutation>();

    Map<String, List<Mutation>> mutationMap = new HashMap<String, List<Mutation>>();

    Map<ByteBuffer, Map<String, List<Mutation>>> record = new HashMap<ByteBuffer, Map<String, List<Mutation>>>();

    ColumnParent parent;

    ConsistencyLevel readConsistencyLevel = ConsistencyLevel.ONE;

    ConsistencyLevel writeConsistencyLevel = ConsistencyLevel.ONE;

    ConsistencyLevel scanConsistencyLevel = ConsistencyLevel.ONE;

    ConsistencyLevel deleteConsistencyLevel = ConsistencyLevel.ONE;

    /**
     * Initialize any state for this DB. Called once per DB instance; there is
     * one DB instance per client thread.
     */
    public void init() throws DBException
    {
//        System.out.println("**************************" + column_family);
        try
        {
            String hosts = getProperties().getProperty("hosts");
            if (hosts == null)
            {
                throw new DBException("Required property \"hosts\" missing for CassandraClient");
            }

            column_family = getProperties().getProperty(COLUMN_FAMILY_PROPERTY);
            parent = new ColumnParent(column_family);

            _table = getProperties().getProperty("schema");
            ConnectionRetries = Integer.parseInt(getProperties().getProperty(CONNECTION_RETRY_PROPERTY,
                    CONNECTION_RETRY_PROPERTY_DEFAULT));

            String username = getProperties().getProperty(USERNAME_PROPERTY);
            String password = getProperties().getProperty(PASSWORD_PROPERTY);

            readConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
                    READ_CONSISTENCY_LEVEL_PROPERTY, READ_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
            writeConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
                    WRITE_CONSISTENCY_LEVEL_PROPERTY, WRITE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
            scanConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
                    SCAN_CONSISTENCY_LEVEL_PROPERTY, SCAN_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));
            deleteConsistencyLevel = ConsistencyLevel.valueOf(getProperties().getProperty(
                    DELETE_CONSISTENCY_LEVEL_PROPERTY, DELETE_CONSISTENCY_LEVEL_PROPERTY_DEFAULT));

            _debug = Boolean.parseBoolean(getProperties().getProperty("debug", "false"));

            String[] allhosts = hosts.split(",");
            String myhost = allhosts[random.nextInt(allhosts.length)];

            Exception connectexception = null;

//            System.out.println("Initiating connection");
            if (tr == null || client == null)
            {

                for (int retry = 0; retry < ConnectionRetries; retry++)
                {
                    tr = new TFramedTransport(new TSocket(myhost, 9160));
                    TProtocol proto = new TBinaryProtocol(tr);
//                    System.out.println("creating connection::");
                    client = new Cassandra.Client(proto);
                    try
                    {
                        tr.open();
                        connectexception = null;
                        break;
                    }
                    catch (Exception e)
                    {
                        connectexception = e;
                        logger.error(e);
                    }
                    try
                    {
                        Thread.sleep(1000);
                    }
                    catch (InterruptedException e)
                    {
                        logger.error(e);
                    }

                }
            }

            if (connectexception != null)
            {
                logger.info("Unable to connect to " + myhost + " after " + ConnectionRetries + " tries");
                throw new DBException(connectexception);
            }

            if (username != null && password != null)
            {
                Map<String, String> cred = new HashMap<String, String>();
                cred.put("username", username);
                cred.put("password", password);
                AuthenticationRequest req = new AuthenticationRequest(cred);
                try
                {
                    client.login(req);
                }
                catch (Exception e)
                {
                    e.printStackTrace();
                    logger.info(e);
                    throw new DBException(e);
                }
            }
            client.set_keyspace(_table);
        }
        catch (Exception e)
        {
            logger.info(e);

        }
    }

    /**
     * Cleanup any state for this DB. Called once per DB instance; there is one
     * DB instance per client thread.
     */
    public void cleanup() throws DBException
    {
        // tr.close();
    }

    /**
     * Read a record from the database. Each field/value pair from the result
     * will be stored in a HashMap.
     * 
     * @param table
     *            The name of the table
     * @param key
     *            The record key of the record to read.
     * @param fields
     *            The list of fields to read, or null for all of them
     * @param result
     *            A HashMap of field/value pairs for the result
     * @return Zero on success, a non-zero error code on error
     */
    public int read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result)
    {
        // System.out.println(key);
//        System.out.println("key:" + key);
//        System.out.println("_table" + _table);
//        System.out.println("table" + table);

        /*
         * if (!_table.equals(table)) {
         */
        try
        {
            /*
             * _table = table;
             */}
        catch (Exception e)
        {
            logger.error(e);
            return Error;
        }
        // }

        try
        {
            SlicePredicate predicate = new SlicePredicate();
            if (fields == null)
            {
//                SlicePredicate predicate = new SlicePredicate();
                predicate = new SlicePredicate();
                predicate.setSlice_range(new SliceRange(Bytes.EMPTY.getBytes(), Bytes.EMPTY.getBytes(), true, 10000));

//                predicate = new SlicePredicate().setSlice_range(new SliceRange(emptyByteBuffer, emptyByteBuffer, false,
//                        1000000));

            }
            else
            {
                ArrayList<ByteBuffer> fieldlist = new ArrayList<ByteBuffer>(fields.size());
                for (String s : fields)
                {
                    fieldlist.add(ByteBuffer.wrap(s.getBytes("UTF-8")));
                }

                predicate = new SlicePredicate().setColumn_names(fieldlist);
            }

            List<ColumnOrSuperColumn> results = client.get_slice(ByteBuffer.wrap(key.getBytes("UTF-8")), parent,
                    predicate, readConsistencyLevel);

            if (_debug)
            {
                System.out.print("Reading key: " + key);
            }

            Column column;
            String name;
            ByteIterator value;
            
            assert results != null && !results.isEmpty();
/*            for (ColumnOrSuperColumn oneresult : results)
            {

                column = oneresult.column;
                name = new String(column.name.array(), column.name.position() + column.name.arrayOffset(),
                        column.name.remaining());
                value = new ByteArrayByteIterator(column.value.array(), column.value.position()
                        + column.value.arrayOffset(), column.value.remaining());

                result.put(name, value);

                if (_debug)
                {
                    logger.info("(" + name + "=" + value + ")");
                }
            }
*/
            if (_debug)
            {
                logger.info("ConsistencyLevel=" + readConsistencyLevel.toString());
            }

            return Ok;
        }
        catch (Exception e)
        {
            logger.error(e);
            return Error;

        }
    }

    /**
     * Perform a range scan for a set of records in the database. Each
     * field/value pair from the result will be stored in a HashMap.
     * 
     * @param table
     *            The name of the table
     * @param startkey
     *            The record key of the first record to read.
     * @param recordcount
     *            The number of records to read
     * @param fields
     *            The list of fields to read, or null for all of them
     * @param result
     *            A Vector of HashMaps, where each HashMap is a set field/value
     *            pairs for one record
     * @return Zero on success, a non-zero error code on error
     */
    public int scan(String table, String startkey, int recordcount, Set<String> fields,
            Vector<HashMap<String, ByteIterator>> result)
    {
        if (!_table.equals(table))
        {
            try
            {
                client.set_keyspace(table);
                _table = table;
            }
            catch (Exception e)
            {
                logger.error(e);
                return Error;
            }
        }
        try
        {
            SlicePredicate predicate;
            if (fields == null)
            {
                predicate = new SlicePredicate().setSlice_range(new SliceRange(emptyByteBuffer, emptyByteBuffer, false,
                        1000000));

            }
            else
            {
                ArrayList<ByteBuffer> fieldlist = new ArrayList<ByteBuffer>(fields.size());
                for (String s : fields)
                {
                    fieldlist.add(ByteBuffer.wrap(s.getBytes("UTF-8")));
                }

                predicate = new SlicePredicate().setColumn_names(fieldlist);
            }

            KeyRange kr = new KeyRange().setStart_key(startkey.getBytes("UTF-8")).setEnd_key(new byte[] {})
                    .setCount(recordcount);

            List<KeySlice> results = client.get_range_slices(parent, predicate, kr, scanConsistencyLevel);

            if (_debug)
            {
                logger.info("Scanning startkey: " + startkey);
            }

            HashMap<String, ByteIterator> tuple;
            for (KeySlice oneresult : results)
            {
                tuple = new HashMap<String, ByteIterator>();

                Column column;
                String name;
                ByteIterator value;
                for (ColumnOrSuperColumn onecol : oneresult.columns)
                {
                    column = onecol.column;
                    name = new String(column.name.array(), column.name.position() + column.name.arrayOffset(),
                            column.name.remaining());
                    value = new ByteArrayByteIterator(column.value.array(), column.value.position()
                            + column.value.arrayOffset(), column.value.remaining());

                    tuple.put(name, value);

                    if (_debug)
                    {
                        logger.info("(" + name + "=" + value + ")");
                    }
                }

                result.add(tuple);
                if (_debug)
                {
                    logger.info("ConsistencyLevel=" + scanConsistencyLevel.toString());
                }
            }

            return Ok;
        }
        catch (Exception e)
        {
            logger.error(e);
            return Error;
        }

    }

    /**
     * Update a record in the database. Any field/value pairs in the specified
     * values HashMap will be written into the record with the specified record
     * key, overwriting any existing values with the same field name.
     * 
     * @param table
     *            The name of the table
     * @param key
     *            The record key of the record to write.
     * @param values
     *            A HashMap of field/value pairs to update in the record
     * @return Zero on success, a non-zero error code on error
     */
    public int update(String table, String key, HashMap<String, ByteIterator> values)
    {
        return insert(table, key, values);
//        return Error;
    }

    /**
     * Insert a record in the database. Any field/value pairs in the specified
     * values HashMap will be written into the record with the specified record
     * key.
     * 
     * @param table
     *            The name of the table
     * @param key
     *            The record key of the record to insert.
     * @param values
     *            A HashMap of field/value pairs to insert in the record
     * @return Zero on success, a non-zero error code on error
     */
    public int insert(String table, String key, HashMap<String, ByteIterator> values)
    {

        /*
         * if (!_table.equals(table)) {
         */try
        {
            client.set_keyspace(_table);
            // _table = table;
        }
        catch (Exception e)
        {
            logger.error(e);
            return Error;
        }
        /*
         * }
         */if (_debug)
        {
//            System.out.println("Inserting key: " + key);
        }

        try
        {
            ByteBuffer wrappedKey = ByteBuffer.wrap(key.getBytes("UTF-8"));

            Column col;
            ColumnOrSuperColumn column;
            for (Map.Entry<String, ByteIterator> entry : values.entrySet())
            {
                col = new Column();
                col.setName(ByteBuffer.wrap(entry.getKey().getBytes("UTF-8")));
                col.setValue(ByteBuffer.wrap(entry.getValue().toArray()));
                col.setTimestamp(System.currentTimeMillis());

                column = new ColumnOrSuperColumn();
                column.setColumn(col);

                mutations.add(new Mutation().setColumn_or_supercolumn(column));
            }

            mutationMap.put(column_family, mutations);
            record.put(wrappedKey, mutationMap);

            client.batch_mutate(record, writeConsistencyLevel);

            mutations.clear();
            mutationMap.clear();
            record.clear();

            if (_debug)
            {
                logger.info("ConsistencyLevel=" + writeConsistencyLevel.toString());
            }

            return Ok;
        }
        catch (Exception e)
        {
            logger.error(e);
            return Error;
        }
    }

    /**
     * Delete a record from the database.
     * 
     * @param table
     *            The name of the table
     * @param key
     *            The record key of the record to delete.
     * @return Zero on success, a non-zero error code on error
     */
    public int delete(String table, String key)
    {
        if (!_table.equals(table))
        {
            try
            {
                client.set_keyspace(table);
                _table = table;
            }
            catch (Exception e)
            {
                logger.error(e);
                return Error;
            }
        }
        try
        {
            client.remove(ByteBuffer.wrap(key.getBytes("UTF-8")), new ColumnPath(column_family),
                    System.currentTimeMillis(), deleteConsistencyLevel);

            if (_debug)
            {
                logger.info("Delete key: " + key);
                logger.info("ConsistencyLevel=" + deleteConsistencyLevel.toString());
            }

            return Ok;
        }
        catch (Exception e)
        {
            logger.error(e);
            return Error;
        }
    }

    public static void main(String[] args)
    {
        ThriftClient cli = new ThriftClient();

        Properties props = new Properties();

        props.setProperty("hosts", "localhost");
        cli.setProperties(props);

        try
        {
            cli.init();
        }
        catch (Exception e)
        {
            logger.error(e);
            System.exit(0);
        }

        HashMap<String, ByteIterator> vals = new HashMap<String, ByteIterator>();
        vals.put("age", new StringByteIterator("57"));
        vals.put("middlename", new StringByteIterator("bradley"));
        vals.put("favoritecolor", new StringByteIterator("blue"));
        int res = cli.insert("usertable", "BrianFrankCooper", vals);
        logger.info("Result of insert: " + res);

        HashMap<String, ByteIterator> result = new HashMap<String, ByteIterator>();
        HashSet<String> fields = new HashSet<String>();
        fields.add("middlename");
        fields.add("age");
        fields.add("favoritecolor");
        res = cli.read("usertable", "BrianFrankCooper", null, result);
        logger.info("Result of read: " + res);
        for (String s : result.keySet())
        {
            logger.info("[" + s + "]=[" + result.get(s) + "]");
        }

        res = cli.delete("usertable", "BrianFrankCooper");
        logger.info("Result of delete: " + res);
    }
}
